"""
速率限制模块 - 防止恶意高频请求

这个模块实现了多种速率限制算法：
- 令牌桶算法
- 滑动窗口算法
- 固定窗口算法
- 分布式速率限制
- 基于用户和IP的限制
"""

import asyncio
import logging
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone, timedelta
from enum import Enum
from typing import Any, Dict, List, Optional, Union, Tuple
import json
import hashlib
from collections import defaultdict, deque

logger = logging.getLogger(__name__)


class RateLimitAlgorithm(Enum):
    """速率限制算法"""
    TOKEN_BUCKET = "token_bucket"
    SLIDING_WINDOW = "sliding_window"
    FIXED_WINDOW = "fixed_window"
    LEAKY_BUCKET = "leaky_bucket"


class RateLimitScope(Enum):
    """速率限制范围"""
    GLOBAL = "global"          # 全局限制
    IP = "ip"                  # IP限制
    USER = "user"              # 用户限制
    API_KEY = "api_key"        # API密钥限制
    ENDPOINT = "endpoint"      # 端点限制


class RateLimitResult(Enum):
    """速率限制结果"""
    ALLOWED = "allowed"        # 允许
    RATE_LIMITED = "rate_limited"  # 被限制
    QUOTA_EXCEEDED = "quota_exceeded"  # 配额超限


@dataclass
class RateLimitConfig:
    """速率限制配置"""
    # 基本配置
    requests_per_minute: int = 60
    requests_per_hour: int = 1000
    requests_per_day: int = 10000
    
    # 算法配置
    algorithm: RateLimitAlgorithm = RateLimitAlgorithm.TOKEN_BUCKET
    window_size: int = 60  # 窗口大小（秒）
    
    # 令牌桶配置
    bucket_capacity: int = 100
    refill_rate: float = 1.0  # 每秒补充令牌数
    
    # 滑动窗口配置
    window_granularity: int = 1  # 窗口粒度（秒）
    
    # 分布式配置
    enable_distributed: bool = False
    redis_url: str = "redis://localhost:6379"
    key_prefix: str = "rate_limit"
    
    # 白名单和黑名单
    whitelist: List[str] = field(default_factory=list)
    blacklist: List[str] = field(default_factory=list)
    
    # 响应配置
    enable_headers: bool = True
    custom_message: str = "Rate limit exceeded"
    
    # 监控配置
    enable_metrics: bool = True
    metrics_retention: int = 86400  # 24小时


@dataclass
class RateLimitInfo:
    """速率限制信息"""
    allowed: bool
    remaining: int
    reset_time: int
    retry_after: Optional[int] = None
    limit: int = 0
    algorithm: Optional[RateLimitAlgorithm] = None


class TokenBucket:
    """令牌桶算法实现"""
    
    def __init__(self, capacity: int, refill_rate: float):
        self.capacity = capacity
        self.refill_rate = refill_rate
        self.tokens = capacity
        self.last_refill = time.time()
        self.lock = asyncio.Lock()
    
    async def consume(self, tokens: int = 1) -> bool:
        """消费令牌"""
        async with self.lock:
            now = time.time()
            # 补充令牌
            time_passed = now - self.last_refill
            tokens_to_add = time_passed * self.refill_rate
            self.tokens = min(self.capacity, self.tokens + tokens_to_add)
            self.last_refill = now
            
            # 检查是否有足够的令牌
            if self.tokens >= tokens:
                self.tokens -= tokens
                return True
            return False
    
    async def get_tokens(self) -> int:
        """获取当前令牌数"""
        async with self.lock:
            now = time.time()
            time_passed = now - self.last_refill
            tokens_to_add = time_passed * self.refill_rate
            self.tokens = min(self.capacity, self.tokens + tokens_to_add)
            self.last_refill = now
            return int(self.tokens)


class SlidingWindow:
    """滑动窗口算法实现"""
    
    def __init__(self, window_size: int, max_requests: int):
        self.window_size = window_size
        self.max_requests = max_requests
        self.requests = deque()
        self.lock = asyncio.Lock()
    
    async def is_allowed(self) -> bool:
        """检查是否允许请求"""
        async with self.lock:
            now = time.time()
            window_start = now - self.window_size
            
            # 移除窗口外的请求
            while self.requests and self.requests[0] < window_start:
                self.requests.popleft()
            
            # 检查是否超过限制
            if len(self.requests) < self.max_requests:
                self.requests.append(now)
                return True
            return False
    
    async def get_remaining(self) -> int:
        """获取剩余请求数"""
        async with self.lock:
            now = time.time()
            window_start = now - self.window_size
            
            # 移除窗口外的请求
            while self.requests and self.requests[0] < window_start:
                self.requests.popleft()
            
            return max(0, self.max_requests - len(self.requests))


class FixedWindow:
    """固定窗口算法实现"""
    
    def __init__(self, window_size: int, max_requests: int):
        self.window_size = window_size
        self.max_requests = max_requests
        self.window_start = int(time.time() // window_size) * window_size
        self.request_count = 0
        self.lock = asyncio.Lock()
    
    async def is_allowed(self) -> bool:
        """检查是否允许请求"""
        async with self.lock:
            now = time.time()
            current_window = int(now // self.window_size) * self.window_size
            
            # 检查是否需要重置窗口
            if current_window > self.window_start:
                self.window_start = current_window
                self.request_count = 0
            
            # 检查是否超过限制
            if self.request_count < self.max_requests:
                self.request_count += 1
                return True
            return False
    
    async def get_remaining(self) -> int:
        """获取剩余请求数"""
        async with self.lock:
            now = time.time()
            current_window = int(now // self.window_size) * self.window_size
            
            # 检查是否需要重置窗口
            if current_window > self.window_start:
                return self.max_requests
            
            return max(0, self.max_requests - self.request_count)


class RateLimiter:
    """速率限制器"""
    
    def __init__(self, config: Optional[RateLimitConfig] = None):
        self.config = config or RateLimitConfig()
        
        # 限制器存储
        self.limiters: Dict[str, Union[TokenBucket, SlidingWindow, FixedWindow]] = {}
        
        # 统计信息
        self.metrics: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int))
        
        # 锁
        self.lock = asyncio.Lock()
        
        self.logger = logging.getLogger(__name__)
        self.logger.info("速率限制器已初始化")
    
    def _generate_key(self, 
                     scope: RateLimitScope,
                     identifier: str,
                     endpoint: Optional[str] = None) -> str:
        """生成限制键"""
        key_parts = [scope.value, identifier]
        if endpoint:
            key_parts.append(endpoint)
        return ":".join(key_parts)
    
    def _is_whitelisted(self, identifier: str) -> bool:
        """检查是否在白名单中"""
        return identifier in self.config.whitelist
    
    def _is_blacklisted(self, identifier: str) -> bool:
        """检查是否在黑名单中"""
        return identifier in self.config.blacklist
    
    async def _get_limiter(self, 
                          scope: RateLimitScope,
                          identifier: str,
                          endpoint: Optional[str] = None) -> Union[TokenBucket, SlidingWindow, FixedWindow]:
        """获取或创建限制器"""
        key = self._generate_key(scope, identifier, endpoint)
        
        if key not in self.limiters:
            if self.config.algorithm == RateLimitAlgorithm.TOKEN_BUCKET:
                limiter = TokenBucket(
                    capacity=self.config.bucket_capacity,
                    refill_rate=self.config.refill_rate
                )
            elif self.config.algorithm == RateLimitAlgorithm.SLIDING_WINDOW:
                limiter = SlidingWindow(
                    window_size=self.config.window_size,
                    max_requests=self.config.requests_per_minute
                )
            elif self.config.algorithm == RateLimitAlgorithm.FIXED_WINDOW:
                limiter = FixedWindow(
                    window_size=self.config.window_size,
                    max_requests=self.config.requests_per_minute
                )
            else:
                raise ValueError(f"不支持的算法: {self.config.algorithm}")
            
            self.limiters[key] = limiter
        
        return self.limiters[key]
    
    async def check_rate_limit(self, 
                              scope: RateLimitScope,
                              identifier: str,
                              endpoint: Optional[str] = None,
                              cost: int = 1) -> RateLimitInfo:
        """检查速率限制"""
        # 检查白名单
        if self._is_whitelisted(identifier):
            return RateLimitInfo(
                allowed=True,
                remaining=999999,
                reset_time=int(time.time() + 3600),
                limit=999999
            )
        
        # 检查黑名单
        if self._is_blacklisted(identifier):
            return RateLimitInfo(
                allowed=False,
                remaining=0,
                reset_time=int(time.time() + 3600),
                retry_after=3600,
                limit=0
            )
        
        try:
            limiter = await self._get_limiter(scope, identifier, endpoint)
            
            # 根据算法检查限制
            if self.config.algorithm == RateLimitAlgorithm.TOKEN_BUCKET:
                allowed = await limiter.consume(cost)
                remaining = await limiter.get_tokens()
            elif self.config.algorithm in [RateLimitAlgorithm.SLIDING_WINDOW, RateLimitAlgorithm.FIXED_WINDOW]:
                allowed = await limiter.is_allowed()
                remaining = await limiter.get_remaining()
            else:
                allowed = False
                remaining = 0
            
            # 更新统计信息
            if self.config.enable_metrics:
                await self._update_metrics(scope, identifier, allowed)
            
            # 计算重置时间
            reset_time = int(time.time() + self.config.window_size)
            
            return RateLimitInfo(
                allowed=allowed,
                remaining=remaining,
                reset_time=reset_time,
                retry_after=None if allowed else self.config.window_size,
                limit=self.config.requests_per_minute,
                algorithm=self.config.algorithm
            )
            
        except Exception as e:
            self.logger.error(f"速率限制检查失败: {e}")
            # 发生错误时允许请求通过
            return RateLimitInfo(
                allowed=True,
                remaining=self.config.requests_per_minute,
                reset_time=int(time.time() + self.config.window_size),
                limit=self.config.requests_per_minute
            )
    
    async def _update_metrics(self, 
                             scope: RateLimitScope,
                             identifier: str,
                             allowed: bool) -> None:
        """更新统计信息"""
        key = f"{scope.value}:{identifier}"
        
        if allowed:
            self.metrics[key]["allowed"] += 1
        else:
            self.metrics[key]["blocked"] += 1
        
        self.metrics[key]["total"] += 1
    
    async def get_metrics(self, 
                         scope: Optional[RateLimitScope] = None,
                         identifier: Optional[str] = None) -> Dict[str, Any]:
        """获取统计信息"""
        if scope and identifier:
            key = f"{scope.value}:{identifier}"
            return dict(self.metrics[key])
        
        # 返回所有统计信息
        return {
            key: dict(metrics) 
            for key, metrics in self.metrics.items()
        }
    
    async def reset_metrics(self) -> None:
        """重置统计信息"""
        self.metrics.clear()
        self.logger.info("速率限制统计信息已重置")
    
    async def add_to_whitelist(self, identifier: str) -> None:
        """添加到白名单"""
        if identifier not in self.config.whitelist:
            self.config.whitelist.append(identifier)
            self.logger.info(f"已添加到白名单: {identifier}")
    
    async def remove_from_whitelist(self, identifier: str) -> None:
        """从白名单移除"""
        if identifier in self.config.whitelist:
            self.config.whitelist.remove(identifier)
            self.logger.info(f"已从白名单移除: {identifier}")
    
    async def add_to_blacklist(self, identifier: str) -> None:
        """添加到黑名单"""
        if identifier not in self.config.blacklist:
            self.config.blacklist.append(identifier)
            self.logger.info(f"已添加到黑名单: {identifier}")
    
    async def remove_from_blacklist(self, identifier: str) -> None:
        """从黑名单移除"""
        if identifier in self.config.blacklist:
            self.config.blacklist.remove(identifier)
            self.logger.info(f"已从黑名单移除: {identifier}")
    
    async def cleanup_expired_limiters(self) -> int:
        """清理过期的限制器"""
        # 这里可以实现清理逻辑，比如清理长时间未使用的限制器
        # 简化实现，返回0
        return 0


class DistributedRateLimiter(RateLimiter):
    """分布式速率限制器"""
    
    def __init__(self, 
                 config: Optional[RateLimitConfig] = None,
                 redis_client=None):
        super().__init__(config)
        self.redis_client = redis_client
        self.logger.info("分布式速率限制器已初始化")
    
    async def _get_limiter(self, 
                          scope: RateLimitScope,
                          identifier: str,
                          endpoint: Optional[str] = None) -> Union[TokenBucket, SlidingWindow, FixedWindow]:
        """获取分布式限制器"""
        if not self.redis_client:
            # 如果没有Redis客户端，回退到本地限制器
            return await super()._get_limiter(scope, identifier, endpoint)
        
        # 这里可以实现基于Redis的分布式限制器
        # 简化实现，使用本地限制器
        return await super()._get_limiter(scope, identifier, endpoint)


class RateLimitMiddleware:
    """速率限制中间件"""
    
    def __init__(self, 
                 rate_limiter: RateLimiter,
                 default_scope: RateLimitScope = RateLimitScope.IP):
        self.rate_limiter = rate_limiter
        self.default_scope = default_scope
        self.logger = logging.getLogger(__name__)
    
    async def __call__(self, request, handler):
        """中间件处理函数"""
        # 提取标识符
        identifier = await self._extract_identifier(request)
        
        # 检查速率限制
        rate_limit_info = await self.rate_limiter.check_rate_limit(
            scope=self.default_scope,
            identifier=identifier,
            endpoint=request.path
        )
        
        # 添加响应头
        if self.rate_limiter.config.enable_headers:
            response_headers = {
                "X-RateLimit-Limit": str(rate_limit_info.limit),
                "X-RateLimit-Remaining": str(rate_limit_info.remaining),
                "X-RateLimit-Reset": str(rate_limit_info.reset_time)
            }
            
            if rate_limit_info.retry_after:
                response_headers["Retry-After"] = str(rate_limit_info.retry_after)
        else:
            response_headers = {}
        
        # 检查是否被限制
        if not rate_limit_info.allowed:
            from aiohttp import web
            return web.json_response(
                {"error": self.rate_limiter.config.custom_message},
                status=429,
                headers=response_headers
            )
        
        # 处理请求
        response = await handler(request)
        
        # 添加响应头
        for key, value in response_headers.items():
            response.headers[key] = value
        
        return response
    
    async def _extract_identifier(self, request) -> str:
        """提取标识符"""
        # 从请求中提取IP地址
        forwarded_for = request.headers.get('X-Forwarded-For')
        if forwarded_for:
            return forwarded_for.split(',')[0].strip()
        
        real_ip = request.headers.get('X-Real-IP')
        if real_ip:
            return real_ip
        
        return request.remote


# 全局速率限制器
_global_rate_limiter: Optional[RateLimiter] = None


def get_rate_limiter() -> RateLimiter:
    """获取全局速率限制器"""
    global _global_rate_limiter
    if _global_rate_limiter is None:
        _global_rate_limiter = RateLimiter()
    return _global_rate_limiter


# 装饰器
def rate_limit(scope: RateLimitScope = RateLimitScope.IP,
               requests_per_minute: int = 60,
               algorithm: RateLimitAlgorithm = RateLimitAlgorithm.TOKEN_BUCKET):
    """速率限制装饰器"""
    def decorator(func):
        async def wrapper(*args, **kwargs):
            # 这里可以实现装饰器逻辑
            # 简化实现，直接调用原函数
            return await func(*args, **kwargs)
        return wrapper
    return decorator
